package com.example.batch.core.health;

import org.springframework.boot.actuate.health.AbstractHealthIndicator;
import org.springframework.boot.actuate.health.Health;
import org.springframework.stereotype.Component;

import java.sql.Connection;
import java.sql.DataSource;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

@Component
public class DataSourceHealthIndicator extends AbstractHealthIndicator {
    
    private final DynamicRoutingDataSource dynamicDataSource;
    private final MultiDataSourceProperties properties;
    
    public DataSourceHealthIndicator(
            DynamicRoutingDataSource dynamicDataSource,
            MultiDataSourceProperties properties) {
        this.dynamicDataSource = dynamicDataSource;
        this.properties = properties;
    }
    
    @Override
    protected void doHealthCheck(Health.Builder builder) {
        Map<String, DataSourceHealth> healthDetails = new HashMap<>();
        
        properties.getSources().forEach((name, config) -> {
            try {
                DataSource dataSource = dynamicDataSource.getResolvedDataSources().get(name);
                try (Connection conn = dataSource.getConnection()) {
                    boolean valid = conn.isValid(3000);
                    DataSourceHealth health = new DataSourceHealth(valid, getMetrics(dataSource));
                    healthDetails.put(name, health);
                }
            } catch (Exception e) {
                healthDetails.put(name, new DataSourceHealth(false, Collections.emptyMap()));
            }
        });
        
        boolean allUp = healthDetails.values().stream()
                .allMatch(DataSourceHealth::isUp);
                
        if (allUp) {
            builder.up();
        } else {
            builder.down();
        }
        
        builder.withDetails(healthDetails);
    }
    
    private Map<String, Object> getMetrics(DataSource dataSource) {
        if (dataSource instanceof HikariDataSource) {
            HikariDataSource hikari = (HikariDataSource) dataSource;
            return Map.of(
                "active", hikari.getHikariPoolMXBean().getActiveConnections(),
                "idle", hikari.getHikariPoolMXBean().getIdleConnections(),
                "total", hikari.getHikariPoolMXBean().getTotalConnections()
            );
        }
        return Collections.emptyMap();
    }
} 